
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class Problem2Optimizer:
    #调整河流位置及管道布局的灌溉系统优化模型

    def __init__(self):
        # 文档核心参数
        self.farm_size = 100
        self.min_soil_moisture = 0.22
        self.pipe_cost = lambda L, Q: 50 * (L ** 1.2) + 0.1 * (Q ** 1.5)
        self.tank_cost_per_liter = 5
        self.sprinkler_radius = 15
        self.sprinkler_min_gap = 15
        self.tank_coverage = 15
        self.soil_dry_weight = 1500
        self.data_period = ('2021-07-01', '2021-07-31')

    def generate_july_data(self):

        print("=== 生成7月灌溉数据 ===")
        dates = pd.date_range(self.data_period[0], self.data_period[1], freq='D')
        n_days = len(dates)

        # 模拟土壤湿度
        np.random.seed(42)
        soil_moisture = np.clip(0.23 + np.random.normal(0, 0.03, n_days), 0.19, 0.35)
        soil_data = pd.DataFrame({'Date': dates, '5cm_SM': soil_moisture})

        # 计算需水量
        soil_data['需水量(L)'] = np.where(
            soil_data['5cm_SM'] < self.min_soil_moisture,
            (self.min_soil_moisture - soil_data['5cm_SM']) * self.soil_dry_weight * 10000 / 100,
            0
        )

        print(f"生成{self.data_period[0]}至{self.data_period[1]}共{n_days}天数据")
        print(f"需灌溉天数：{sum(soil_data['需水量(L)'] > 0)}天，总需水量：{soil_data['需水量(L)'].sum():.2f}L")
        return soil_data

    def optimize_sprinklers(self):
        #按作物分区优化喷头布置（减少冗余）
        # 作物分区：大豆（0 - 30m）、玉米（30 - 60m）、高粱（60 - 100m）
        zones = {
            '大豆': {'y_range': (0, 20), 'x_coords': np.arange(15, 100, 30)},
            '玉米': {'y_range': (20, 50), 'x_coords': np.arange(15, 100, 30)},
            '高粱': {'y_range': (50, 100), 'x_coords': np.arange(15, 100, 30)}
        }

        sprinkler_positions = []
        for zone in zones.values():
            y_start, y_end = zone['y_range']
            y_coords = np.arange(y_start + 15, y_end, 30)
            for x in zone['x_coords']:
                for y in y_coords:
                    sprinkler_positions.append((x, y))

        # 验证喷头间距≥15m
        min_gap = min(np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
                      for i, (x1, y1) in enumerate(sprinkler_positions)
                      for j, (x2, y2) in enumerate(sprinkler_positions) if i < j)
        assert min_gap >= self.sprinkler_min_gap, f"喷头间距{min_gap:.1f}m < 15m（违反图2约束）"

        # 喷头成本（单个200元，按分区数量精简）
        n_sprinklers = len(sprinkler_positions)
        sprinkler_cost = n_sprinklers * 200
        print(f"\n喷头布置：{n_sprinklers}个（分区30m网格），最小间距{min_gap:.1f}m，成本{sprinkler_cost:.2f}元")

        # 打印喷头坐标
        print("\n喷头坐标列表（x, y）：")
        for i, (x, y) in enumerate(sprinkler_positions, 1):
            print(f"  喷头{i}: ({x:.1f}m, {y:.1f}m)")

        return n_sprinklers, sprinkler_positions, sprinkler_cost, zones

    def optimize_pipes_and_tanks(self, soil_data, sprinkler_positions, zones):
        #优化分区支管与储水罐（最小化总长度）
        print("\n=== 优化管道与储水罐 ===")
        daily_need = soil_data[soil_data['需水量(L)'] > 0]['需水量(L)']
        avg_daily = daily_need.mean()
        max_daily = daily_need.max()


        pipe_segments = []
        total_pipe_length = 0

        for zone_name, zone in zones.items():
            y_start, y_end = zone['y_range']
            for x, y in sprinkler_positions:
                if y_start <= y < y_end:

                    pipe_length = y
                    pipe_segments.append((pipe_length, avg_daily))
                    total_pipe_length += pipe_length

        # 计算总管道成本
        total_pipe_cost = sum(self.pipe_cost(L, Q) for L, Q in pipe_segments)
        print(f"管道系统：垂直分管道，总长{total_pipe_length:.1f}m，总成本{total_pipe_cost:.2f}元")

        # 打印管道日流量信息
        print(f"\n管道日流量信息：")
        print(f"  平均日流量：{avg_daily:.2f}L/天")
        print(f"  最大日流量：{max_daily:.2f}L/天")
        print(f"  各分段流量：均为平均日流量{avg_daily:.2f}L/天（按分区设计）")

        # 2. 储水罐优化（每个作物区1个，位于分区中心）
        tank_positions = [((zones[name]['y_range'][0] + zones[name]['y_range'][1]) / 2, 50)
                          for name in zones]  # (y, x)
        # 验证储水罐覆盖半径15m内有喷头
        for tx, ty in tank_positions:
            assert any(np.sqrt((x - ty) ** 2 + (y - tx) ** 2) <= self.tank_coverage
                       for x, y in sprinkler_positions), f"储水罐({tx},{ty})未覆盖喷头"

        # 储水罐总容积 = 1.2倍最大日需水量（按分区分配）
        tank_volume = sum(max_daily * 1.2 * (len([p for p in sprinkler_positions if
                                                  zones[name]['y_range'][0] <= p[1] < zones[name]['y_range'][1]]) / len(
            sprinkler_positions))
                          for name in zones)
        tank_cost = tank_volume * self.tank_cost_per_liter
        print(f"储水罐：3个（分区中心），总容积{tank_volume:.2f}L，成本{tank_cost:.2f}元")

        return total_pipe_cost, tank_cost, total_pipe_length, tank_volume, tank_positions, zones

    def draw_layout(self, sprinklers, tanks, zones, pipe_length):
        #绘制调整后的灌溉系统规划图
        fig, ax = plt.subplots(figsize=(12, 10))

        # 1. 农场边界与作物分区
        ax.add_patch(plt.Rectangle((0, 0), 100, 100, fill=False, edgecolor='green', linewidth=3, label='农场边界'))
        for name, zone in zones.items():
            y_start, y_end = zone['y_range']
            ax.add_patch(plt.Rectangle((0, y_start), 100, y_end - y_start,
                                       fill=False, edgecolor='gray', linestyle='--'))
            ax.text(50, (y_start + y_end) / 2, name, ha='center', va='center', fontweight='bold')

        # 2. 河流（农场下方y = -5至0）
        ax.fill_betweenx([-5, 0], [0, 0], [100, 100], color='cyan', alpha=0.5, label='河流')

        # 3. 喷头及覆盖范围
        for x, y in sprinklers:
            ax.scatter(x, y, color='blue', s=80, zorder=5)
            ax.add_patch(plt.Circle((x, y), self.sprinkler_radius, fill=False, color='blue', alpha=0.2))

        # 4. 垂直分管道
        for x, y in sprinklers:
            ax.plot([x, x], [y, 0], color='brown', linewidth=1.5, alpha=0.7, label='垂直分管道' if y == 15 else "")

        # 5. 储水罐及覆盖范围
        for i, (tx, ty) in enumerate(tanks):
            ax.scatter(ty, tx, color='orange', s=150, marker='s', zorder=5)
            ax.add_patch(plt.Circle((ty, tx), self.tank_coverage, fill=False, color='orange', alpha=0.3))
            ax.text(ty, tx - 8, f'储水罐{i + 1}', fontsize=10, ha='center')

        # 标注
        ax.set_xlim(-10, 110)
        ax.set_ylim(-10, 105)
        ax.set_xlabel('距离(米)')
        ax.set_ylabel('距离(米)')
        ax.set_title('调整后的灌溉系统规划图（河流在下方，垂直分管道）')
        ax.legend(loc='upper right')
        ax.grid(alpha=0.3)
        ax.set_aspect('equal')
        plt.tight_layout()
        plt.show()

    def run_optimization(self):
        #执行完整优化流程
        # 步骤1：生成数据
        soil_data = self.generate_july_data()

        # 步骤2：优化喷头
        n_sprinklers, sprinkler_pos, sprinkler_cost, zones = self.optimize_sprinklers()

        # 步骤3：优化管道和储水罐
        pipe_cost, tank_cost, pipe_len, tank_vol, tanks, zones = \
            self.optimize_pipes_and_tanks(soil_data, sprinkler_pos, zones)

        # 步骤4：总费用
        total_cost = pipe_cost + tank_cost + sprinkler_cost
        print(f"\n=== 优化结果 ===")
        print(f"总建设成本：{total_cost:.2f}元（管道+储水罐+喷头）")
        print(f"配置：管道{pipe_len:.1f}m，储水罐{tank_vol:.2f}L，喷头{n_sprinklers}个")

        # 可视化需水量
        plt.figure(figsize=(12, 5))
        plt.bar(soil_data['Date'], soil_data['需水量(L)'], color='skyblue')
        plt.axhline(y=0, color='red', linestyle='--', label=f'最低湿度阈值({self.min_soil_moisture})')
        plt.title('2021年7月每日需水量')
        plt.xlabel('日期')
        plt.ylabel('需水量(L)')
        plt.xticks(rotation=45)
        plt.show()

        # 绘制规划图
        self.draw_layout(sprinkler_pos, tanks, zones, pipe_len)

        return {'总费用': total_cost, '管道成本': pipe_cost, '储水罐成本': tank_cost}


if __name__ == "__main__":
    optimizer = Problem2Optimizer()
    result = optimizer.run_optimization()
